#include <defs.h>

.extern trap_handler
.extern dummy
.extern swapper_pg_dir
.section .text.entry
.align 2
.global _traps
.global __dummy
.global __switch_to

# PUSH宏，用于压栈 
.macro PUSH reg
    sd  \reg, 0(sp)
    addi sp,sp,-8
.endm
# POP宏，用于出栈
.macro POP reg
    addi sp,sp,8
    ld  \reg, 0(sp)
.endm

.macro PRINT str
    la a0, \str
    mv a1, x0
    mv a2, x0
    mv a3, x0
    mv a4, x0
    mv a5, x0
    li a6, 0
    li a7, 1
    ecall
.endm

_traps:
    # YOUR CODE HERE
    # -----------
        # ecall print a char
        # li a6, 0x00
        # li a7, 0x01
        # li a0, 'A'
        # ecall

    _traps_start:
        //If sscratch is not 0, then switch sp and sscratch
        csrrw sp, sscratch, sp


        # 1. save 32 registers and sepc to stack
    _traps_pushR:  
        PUSH x2
        PUSH x1
        PUSH x3
        PUSH x4
        PUSH x5
        PUSH x6
        PUSH x7
        PUSH x8
        PUSH x9
        PUSH x10
        PUSH x11
        PUSH x12
        PUSH x13
        PUSH x14
        PUSH x15
        PUSH x16
        PUSH x17
        PUSH x18
        PUSH x19
        PUSH x20
        PUSH x21
        PUSH x22
        PUSH x23
        PUSH x24
        PUSH x25
        PUSH x26
        PUSH x27
        PUSH x28
        PUSH x29
        PUSH x30
        PUSH x31
        csrr t0, sstatus
        PUSH t0
        csrr t0, scause
        PUSH t0
        csrr t0, stval
        PUSH t0
        csrr t0, sepc
        PUSH t0
        
    # -----------

        # 2. call trap_handler
        csrr a0, scause
        csrr a1, sepc
        mv a2, sp
        addi a2, a2, 8
        call trap_handler


    # -----------

        # 3. restore sepc and 32 registers (x2(sp) should be restore last) from stack

    _traps_popR:
        POP t0
        csrw sepc, t0
        POP t0
        csrw stval, t0
        POP t0
        csrw scause, t0
        POP t0
        csrw sstatus, t0
        POP x31
        POP x30
        POP x29
        POP x28
        POP x27
        POP x26
        POP x25
        POP x24
        POP x23
        POP x22
        POP x21
        POP x20
        POP x19
        POP x18
        POP x17
        POP x16
        POP x15
        POP x14
        POP x13
        POP x12
        POP x11
        POP x10
        POP x9
        POP x8
        POP x7
        POP x6
        POP x5
        POP x4
        POP x3
        POP x1
        POP x2
    # -----------
        # 4. return from trap

    _traps_end:
        # 5. if sscratch is not 0, then switch sp and sscratch
        csrr t0, sscratch
        beq t0, x0, _traps_ret 
        csrw sscratch, sp
        mv sp, t0

    _traps_ret:
        sret

    # -----------
    .global __dummy
# __dummy:
        # la t0, dummy
        # csrw sepc, t0

        //switch sp and sscratch
        # mv t0, sp
        # csrr sp, sscratch
        # csrw sscratch, t0

        # sret

__switch_to:
        sd ra, 40(a0)
        sd sp, 48(a0)
        sd s0, 56(a0)
        sd s1, 64(a0)
        sd s2, 72(a0)
        sd s3, 80(a0)
        sd s4, 88(a0)
        sd s5, 96(a0)
        sd s6, 104(a0)
        sd s7, 112(a0)
        sd s8, 120(a0)
        sd s9, 128(a0)
        sd s10, 136(a0)
        sd s11, 144(a0)

        //store sepc, sstatus, sscratch
        csrr t0,sepc
        sd t0, 152(a0)
        csrr t0,sstatus
        sd t0, 160(a0)
        csrr t0,sscratch
        sd t0, 168(a0)

        //switch the pagetable
        //t1 stores the physical address
        csrr t1, satp


        # li t0, 0xfffffffffff
        # and t1, t1, t0
        # sd t1, 176(a0)

      
        ld ra, 40(a1)
        ld sp, 48(a1)
        ld s0, 56(a1)
        ld s1, 64(a1)
        ld s2, 72(a1)
        ld s3, 80(a1)
        ld s4, 88(a1)
        ld s5, 96(a1)
        ld s6, 104(a1)
        ld s7, 112(a1)
        ld s8, 120(a1)
        ld s9, 128(a1)
        ld s10, 136(a1)
        ld s11, 144(a1)


        //load sepc, sstatus, sscratch
        ld t0, 152(a1)
        csrw sepc,t0
        ld t0, 160(a1)
        csrw sstatus,t0
        ld t0, 168(a1)
        csrw sscratch,t0


        # // mode 8
        # li t0, 0x8000000000000000
        # li t2, PA2VA_OFFSET
        # //switch that to physical address
        # sub t1,t1,t2
        # srli t1, t1, 12
        # or t0, t0, t1

        ld t0, 176(a1)
        //t0 = t0 >> 12
        //srli t0,t0,12
        //set mode 8
        //li t1, 0x8000000000000000
        //or t0,t0,t1
        csrw satp,t0

        sfence.vma zero, zero
        fence.i
        ret        
